import torch
from torch_geometric.data import Dataset, Data
import trimesh
import numpy as np

def simplifyMesh(mesh, simplify_size=417114):
    # print(f"Simplify mesh from {len(mesh.vertices)} to {simplify_size} ")
    simplifed = mesh.simplify_quadric_decimation(face_count=simplify_size)
    return simplifed
    
class STLGraphDataset(Dataset):
    def __init__(self, stl_files, labels, input_dim=6, transform=None, simplify_size=208490):
        self.stl_files = stl_files
        self.labels = labels
        self.transform = transform
        self.input_dim = input_dim
        self.simplify_size = simplify_size
        
    def __len__(self):
        return len(self.stl_files)
   
    def __getitem__(self, idx):
        stl_file = self.stl_files[idx]
        label = self.labels[idx]
        mesh = trimesh.load_mesh(stl_file)
        
        # simplify mesh if too large
        if self.simplify_size and len(mesh.vertices) > 1e6:
            mesh = simplifyMesh(mesh, self.simplify_size)
            
        vertices = mesh.vertices
        faces = mesh.faces
        normals = mesh.vertex_normals
    
        x = torch.tensor(np.concatenate([vertices, normals],-1), dtype=torch.float32)[:, :self.input_dim] 
        edge_index = []
        for face in faces:
            edge_index.append([face[0], face[1]])
            edge_index.append([face[1], face[2]])
            edge_index.append([face[2], face[0]])
        
        edge_index = torch.tensor(edge_index, dtype=torch.long).t().contiguous()
        
        y = torch.tensor([label], dtype=torch.float32)
        batch = torch.tensor([0] * x.shape[0], dtype=torch.long)

        data = Data(x=x, edge_index=edge_index, y=y, batch=batch)
        
        if self.transform:
            data = self.transform(data)
        
        return data
    
class NormalizeAugment:
    def __init__(self, augment=False, mode='standard'):
        self.augment = augment
        self.mode = mode 
        
    def __call__(self, data):
        if self.mode == 'standard':
            x_mean = data.x.mean(dim=0)
            x_std = data.x.std(dim=0)
            data.x = (data.x - x_mean) / x_std
            
        elif self.mode == 'normalize':
            x_min, x_max = torch.aminmax(data.x, dim=0)
            data.x = (data.x - x_min) / (x_max - x_min) 
            
        elif self.mode == 'carLength':
            x_min, x_max = torch.aminmax(data.x[:,:1])
            data.x /= (x_max - x_min)
        
        if self.augment and np.random.rand() < 0.5:
            data.x[:, [1,4]] *= -1.
        return data

    